#ifndef CLMEM_H
#define	CLMEM_H

#include "clUtils.h"
#include <algorithm>
#include "Profiler.h"
#include "ClEvent.h"
#include <map>

static const char fillSource[] = 
"typedef struct _DataType { \n"
"	char bytes[%d]; \n"
"} DataType; \n"
" \n"
"__kernel void fill(__global DataType *v, int offset, DataType value) { \n"
"	uint i = offset + get_global_id(0); \n"
"	v[i] = value; \n"
"} \n";

static std::map<size_t, cl_kernel> hashToKernel;

template <typename T>
struct ClMem {
		
	ClMem() : handle(NULL_MEM) { }

	ClMem(cl_command_queue _queue, size_t count) : queue(_queue) {
		Measure m("ClMem()");
		cl_int errorCode;
		handle = clCreateBuffer(getContext(queue), CL_MEM_READ_WRITE, count * sizeof(T), NULL, &errorCode);
		CL_CHECK(errorCode);
	}
	
	ClMem(cl_command_queue _queue, size_t count, const T &value)  : queue(_queue) {
		cl_int errorCode;
		handle = clCreateBuffer(getContext(queue), CL_MEM_READ_WRITE, count * sizeof(T), NULL, &errorCode);
		CL_CHECK(errorCode);
		beginFill(value, 0, count).wait();
	}
	
	ClMem(cl_command_queue _queue, size_t count, const T* data) : queue(_queue) {
		Measure m("ClMem()+init");
		cl_int errorCode;
		handle = clCreateBuffer(getContext(queue), 
				CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR, 
				count * sizeof(T), 
				(void*)data, &errorCode);
		CL_CHECK(errorCode);
	}
		
	ClMem(const ClMem<T> &other) : queue(other.queue), handle(other.handle) {
		if (handle != NULL_MEM) {
			CL_CHECK(clRetainMemObject(handle));
		}
	}
		
	ClMem& operator = (const ClMem<T> &other) {
		ClMem<T> temp(other);
		std::swap(queue, temp.queue);
		std::swap(handle, temp.handle);
		return *this;
	}
	
	inline operator cl_mem () const {
        return handle;
    }
    
    size_t count() const {
        size_t numBytes;
        CL_CHECK(clGetMemObjectInfo(handle, CL_MEM_SIZE, sizeof(size_t), &numBytes, NULL));
        return numBytes / sizeof(T);
    }
	
	ClEvent beginWrite(const T* src, size_t memOffset, size_t count) {
		if (handle == NULL_MEM) throw Exception("Memory handle is null");

		cl_event event;
		CL_CHECK(clEnqueueWriteBuffer(queue, handle, CL_FALSE, 
			memOffset * sizeof(T), 
			count * sizeof(T),  
			src, 0, NULL, &event));
		return ClEvent(event);
	}
	
	ClEvent beginRead(T* dest, size_t memOffset, size_t count) {
		if (handle == NULL_MEM) throw Exception("Memory handle is null");

		cl_event event;
		CL_CHECK(clEnqueueReadBuffer(queue, handle, CL_FALSE, 
			memOffset * sizeof(T), 
			count * sizeof(T), 
			dest, 0, NULL, &event));
		return ClEvent(event);
	}
	
	ClEvent beginFill(const T &value, size_t offset, size_t count) {
		cl_kernel fillKernel = getFillKernel(*this);
		CL_CHECK(clSetKernelArg(fillKernel, 0, sizeof(handle), &handle));
		CL_CHECK(clSetKernelArg(fillKernel, 1, sizeof(offset), &offset));
		CL_CHECK(clSetKernelArg(fillKernel, 2, sizeof(T), &value));
		
		count = count - offset;
		cl_event event;
		CL_CHECK(clEnqueueNDRangeKernel(queue,
				fillKernel, 1, NULL,
				&count, NULL,
				0, NULL, &event));
		return ClEvent(event);
	}

	void write(const T* src, size_t memOffset, size_t count) {
		beginWrite(src, memOffset, count).wait();
	}

	void read(T* dest, size_t memOffset, size_t count) {
		beginRead(dest, memOffset, count).wait();
	}
	
	int refCount() {
		if (handle == NULL_MEM) throw Exception("Memory handle is null");

		int count;
		CL_CHECK(clGetMemObjectInfo(handle, CL_MEM_REFERENCE_COUNT, sizeof(int), &count, NULL));
		return count;
	}
	
	~ClMem() {
		if (handle != NULL_MEM) {
			clReleaseMemObject(handle);
		}
	}
private:
	static const cl_mem NULL_MEM;	
	
	cl_command_queue queue;
	cl_mem handle;
	
	static size_t getHash(cl_context context, cl_device_id device, size_t elementSize) {
		return (size_t)context ^ (size_t)device + elementSize;
	}
	
	static cl_kernel buildKernel(cl_context context, cl_device_id device, size_t elementSize) {
		char src[sizeof(fillSource) / sizeof(char) + 10];
		sprintf(src, fillSource, elementSize);
		cl_program program = buildProgram(context, device, src);
		cl_int errorCode;	
		cl_kernel res = clCreateKernel(program, "fill", &errorCode);
		CL_CHECK(errorCode);
		return res;
	}
	
	template <class U>
	static cl_kernel getFillKernel(const ClMem<U> &mem) {
		cl_context context = getContext(mem.queue);
		cl_device_id deviceId = getDeviceId(mem.queue);
		
		size_t key = getHash(context, deviceId, sizeof(U));
		auto it = hashToKernel.find(key);
		if (it != hashToKernel.end()) return it->second;
		
		cl_kernel res = buildKernel(context, deviceId, sizeof(U));
		hashToKernel.insert(std::pair<size_t, cl_kernel>(key, res));
		return res;
	}
};

template <typename T>
const cl_mem ClMem<T>::NULL_MEM = cl_mem();

#endif	/* CLMEM_H */

